import numpy as np
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit

# 设置字体为 Times New Roman
plt.rcParams['font.family'] = 'Times New Roman'

# 读取数据
data = pd.read_excel('./assets/normalized_data.xlsx')

# 获取年份、人数和税收数据
years = data['Years'].astype(int)
visitor = pd.to_numeric(data['Visitor'], errors='coerce').values
tax = pd.to_numeric(data['Tax'], errors='coerce').values

# 删除 NaN 值
valid_indices = ~np.isnan(visitor) & ~np.isnan(tax)
valid_years = years[valid_indices]
valid_visitors = visitor[valid_indices]
valid_tax = tax[valid_indices]

# 定义税收模型函数
def tax_model(t, a, b, c):
    N_t = np.interp(t, valid_years, valid_visitors)  # 插值后的游客数量
    return c + b * N_t + a * t

# 提供初始参数估计值
initial_params = [1, 1, 1]

# 拟合数据
popt, pcov = curve_fit(tax_model, valid_years, valid_tax, p0=initial_params)
a, b, c = popt

# 打印拟合参数
print(f"a: {a:.6f}")
print(f"b: {b:.6f}")
print(f"c: {c:.6f}")

# 使用 seaborn 绘图
sns.set_theme(style="darkgrid")
plt.figure(figsize=(10, 6))
sns.scatterplot(x=valid_years, y=valid_tax / 1e6, color=sns.color_palette("PuBu")[2], label='Actual Data')
sns.lineplot(x=valid_years, y=tax_model(valid_years, *popt) / 1e6, color=sns.color_palette("PuBu")[4], label='Fitted Curve')
plt.xlabel('Years', fontsize=14, fontname='Times New Roman')
plt.ylabel('Tax (in millions)', fontsize=14, fontname='Times New Roman')
plt.title('Tax Model Fit', fontsize=16, fontname='Times New Roman')
plt.legend()
plt.xticks(ticks=np.arange(valid_years.min(), valid_years.max() + 1, 1), fontname='Times New Roman')
plt.yticks(fontname='Times New Roman')
plt.show()